import numpy as np
import netCDF4 as nc
import os
from datetime import datetime

time = datetime(
    year = 2025, 
    month = 1, 
    day = 1, 
    hour = 0, 
    minute = 0, 
)

forecast_dir = os.path.join(os.getcwd(), 'forecasts', time.strftime('%Y-%m-%d-%H-%M'))
input_dir = os.path.join(os.getcwd(), 'input', time.strftime('%Y-%m-%d-%H-%M'))
os.makedirs(input_dir, exist_ok=True)

# Convert the surface data to npy
surface_data = np.zeros((4, 721, 1440), dtype=np.float32)
with nc.Dataset(os.path.join(forecast_dir , 'surface.nc')) as nc_file:
    surface_data[0] = nc_file.variables['msl'][:].astype(np.float32)
    surface_data[1] = nc_file.variables['u10'][:].astype(np.float32)
    surface_data[2] = nc_file.variables['v10'][:].astype(np.float32)
    surface_data[3] = nc_file.variables['t2m'][:].astype(np.float32)
np.save(os.path.join(input_dir, 'input_surface.npy'), surface_data)

# Convert the upper air data to npy
upper_data = np.zeros((5, 13, 721, 1440), dtype=np.float32)
with nc.Dataset(os.path.join(forecast_dir , 'upper.nc')) as nc_file:
    upper_data[0] = nc_file.variables['z'][:].astype(np.float32)
    upper_data[1] = nc_file.variables['q'][:].astype(np.float32)
    upper_data[2] = nc_file.variables['t'][:].astype(np.float32)
    upper_data[3] = nc_file.variables['u'][:].astype(np.float32)
    upper_data[4] = nc_file.variables['v'][:].astype(np.float32)
np.save(os.path.join(input_dir, 'input_upper.npy'), upper_data)


